package com.yhq.yhqproject.base.modules.chat.service;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.yhq.yhqproject.base.common.web.ResultVO;
import com.yhq.yhqproject.base.modules.chat.entity.GptConfigInfo;
import com.yhq.yhqproject.base.modules.chat.entity.OpenAiChatCompletionParam;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

@Service
@Slf4j
public class ChatGPTService {


    private static final String GPT_STOP = "stop";
    private static final String GPT_CONTINUE = "继续";
    private final GptConfigInfo gptConfigInfo;
    private final RestTemplate restTemplate;

    public ChatGPTService(GptConfigInfo gptConfigInfo, RestTemplate restTemplate) {
        this.gptConfigInfo = gptConfigInfo;
        this.restTemplate = restTemplate;
    }

    /**
     * 这段代码用于向 OpenAI 发送请求并获取回复
     */
    public String requestChatForResult(String prompt) {
        // 创建 OpenAiChatCompletionParam 对象并设置参数
        OpenAiChatCompletionParam param = new OpenAiChatCompletionParam();
        // 设置模型名称
        param.setModel(gptConfigInfo.getOpenAiParamModel());
        //param.setMax_tokens(gptConfigInfo.getOpenAiParamMaxTokens());
        param.setTemperature(gptConfigInfo.getOpenAiParamTemperature());
        param.setTop_p(gptConfigInfo.getOpenAiParamTopP());
        List<JSONObject> messages = new ArrayList<>();
        // 创建用于存储用户聊天记录的 JSONArray
        JSONObject item = new JSONObject();
        item.put("role", "user");
        item.put("content", prompt);
        messages.add(item);
        param.setMessages(messages);
        // 创建用于存储生成的聊天记录的 StringBuilder
        StringBuilder promptResult = new StringBuilder();
        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        HttpEntity<OpenAiChatCompletionParam> httpEntity = new HttpEntity<>(param, headers);
        int limit = 0;
        while (true) {
            if (limit >= gptConfigInfo.getOpenAiChatRetryLimit()) {
                break;
            }
            try {
                log.info("ChatGPT param:{},", JSONObject.toJSONString(param));

                HttpComponentsClientHttpRequestFactory factory = new HttpComponentsClientHttpRequestFactory();
                factory.setConnectTimeout(5 * 60 * 1000); // 连接超时时间为 5 分钟
                factory.setReadTimeout(5 * 60 * 1000); // 读取超时时间为 5 分钟
                restTemplate.setRequestFactory(factory);

                ResultVO result = restTemplate.postForObject(gptConfigInfo.getChatCompletionsUrl(), httpEntity, ResultVO.class);
                log.info("ChatGPT result:{},", JSONObject.toJSONString(result));
                if (result == null || result.getCode() != 0) {
                    //log.error("gpt request error, param:{}, result={}, ", JSONObject.toJSONString(param), JSONObject.toJSONString(result));
                    throw new RuntimeException("gpt request error");
                }
                JSONObject resultData = JSONObject.parseObject(JSONObject.toJSONString(result.getData()));
                JSONArray choices = resultData.getJSONArray("choices");
                JSONObject message = choices.getJSONObject(0).getJSONObject("message");
                String finishReason = choices.getJSONObject(0).getString("finish_reason");
                String finishDetails = message.getString("finish_details");
                String content = message.getString("content");
                promptResult.append(content);

                // 如果stop 或者 finish_reason 为stop 则结束
                if (GPT_STOP.equals(finishDetails) || GPT_STOP.equals(finishReason)) {
                    break;
                } else {
                    item = new JSONObject();
                    item.put("role", "assistant");
                    item.put("content", content);
                    messages.add(item);
                    item = new JSONObject();
                    item.put("role", "user");
                    item.put("content", GPT_CONTINUE);
                    messages.add(item);
                    param.setMessages(messages);
                    httpEntity = new HttpEntity<>(param, headers);
                    TimeUnit.SECONDS.sleep(3);
                }
            } catch (Exception e) {
                log.error("gpt request error, limit={}, prompt={}, ", limit, prompt, e);
                limit++;
                try {
                    TimeUnit.SECONDS.sleep(5);
                } catch (InterruptedException ex) {
                    ex.printStackTrace();
                }
                item = new JSONObject();
                item.put("role", "user");
                item.put("content", prompt);
                messages = new ArrayList<>();
                messages.add(item);
                param.setMessages(messages);
                httpEntity = new HttpEntity<>(param, headers);
            }
        }
        String resp = promptResult.toString();
        if (resp.contains("\uD83D\uDC4D")) {
            resp = resp.replace("\uD83D\uDC4D", "").trim();
        }
        return resp;
    }

}
